//go:build sql_integration

package requestmgr

import (
	"context"
	"fmt"
	"strings"
	"testing"

	imageCVEDSMocks "github.com/stackrox/rox/central/cve/image/v2/datastore/mocks"
	imageDSMocks "github.com/stackrox/rox/central/image/datastore/mocks"
	imageV2DSMocks "github.com/stackrox/rox/central/imagev2/datastore/mocks"
	reprocessorMocks "github.com/stackrox/rox/central/reprocessor/mocks"
	sensorConnMgrMocks "github.com/stackrox/rox/central/sensor/service/connection/mocks"
	vulnReqCache "github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/cache"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/common"
	vulReqDS "github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/datastore"
	dsMock "github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/datastore/mocks"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/utils"
	"github.com/stackrox/rox/generated/storage"
	"github.com/stackrox/rox/pkg/features"
	"github.com/stackrox/rox/pkg/fixtures"
	"github.com/stackrox/rox/pkg/grpc/authn"
	mockIdentity "github.com/stackrox/rox/pkg/grpc/authn/mocks"
	"github.com/stackrox/rox/pkg/postgres/pgtest"
	"github.com/stackrox/rox/pkg/protocompat"
	"github.com/stackrox/rox/pkg/sac"
	"github.com/stackrox/rox/pkg/search"
	"github.com/stretchr/testify/assert"
	"go.uber.org/mock/gomock"
)

func TestCreateWithUnifiedDeferral(t *testing.T) {
	testCreate(t)
}

func testCreate(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	datastore := dsMock.NewMockDataStore(mockCtrl)
	manager := New(nil, datastore, vulnReqCache.New(), nil, nil, nil, nil, nil, nil, nil)

	cve := "CVE-2021-1031"
	globalDefReq := fixtures.GetGlobalDeferralRequestV2(cve)
	globalFPReq := fixtures.GetGlobalFPRequestV2(cve)

	globalDefReqApproved := fixtures.GetGlobalDeferralRequest(cve)
	globalDefReqApproved.Status = storage.RequestStatus_APPROVED
	globalFPReqApproved := fixtures.GetGlobalDeferralRequest(cve)
	globalFPReqApproved.Status = storage.RequestStatus_APPROVED

	allAllowedCtx := sac.WithAllAccess(context.Background())

	cases := []struct {
		name         string
		existingReqs []*storage.VulnerabilityRequest
		reqType      storage.VulnerabilityState
		allow        bool
	}{
		{
			name:         "[DEFER] allow if no other reqs for this cve",
			existingReqs: []*storage.VulnerabilityRequest{},
			reqType:      storage.VulnerabilityState_DEFERRED,
			allow:        true,
		},
		{
			name:         "[DEFER] allow if scope is different",
			existingReqs: []*storage.VulnerabilityRequest{fixtures.GetImageScopeDeferralRequest("registry.k8s.io", "kube-proxy", ".*", cve)},
			reqType:      storage.VulnerabilityState_DEFERRED,
			allow:        true,
		},
		{
			name:         "[DEFER] disallow if exact same cve + scope pending request already exists",
			existingReqs: []*storage.VulnerabilityRequest{globalDefReq},
			reqType:      storage.VulnerabilityState_DEFERRED,
			allow:        false,
		},
		{
			name:         "[DEFER] disallow if exact same cve + scope approved request already exists",
			existingReqs: []*storage.VulnerabilityRequest{globalDefReqApproved},
			reqType:      storage.VulnerabilityState_DEFERRED,
			allow:        false,
		},
		{
			name:         "[DEFER] disallow if exact same cve + scope request already exists even if for FP",
			existingReqs: []*storage.VulnerabilityRequest{globalFPReq},
			reqType:      storage.VulnerabilityState_DEFERRED,
			allow:        false,
		},
		{
			name:         "[FP] allow if no other reqs for this cve",
			existingReqs: []*storage.VulnerabilityRequest{},
			reqType:      storage.VulnerabilityState_FALSE_POSITIVE,
			allow:        true,
		},
		{
			name:         "[FP] allow if scope is different",
			existingReqs: []*storage.VulnerabilityRequest{fixtures.GetImageScopeFPRequest("registry.k8s.io", "kube-proxy", ".*", cve)},
			reqType:      storage.VulnerabilityState_FALSE_POSITIVE,
			allow:        true,
		},
		{
			name:         "[FP] disallow if exact same cve + scope pending request already exists",
			existingReqs: []*storage.VulnerabilityRequest{globalFPReq},
			reqType:      storage.VulnerabilityState_FALSE_POSITIVE,
			allow:        false,
		},
		{
			name:         "[FP] disallow if exact same cve + scope approved request already exists",
			existingReqs: []*storage.VulnerabilityRequest{globalFPReqApproved},
			reqType:      storage.VulnerabilityState_FALSE_POSITIVE,
			allow:        false,
		},
		{
			name:         "[FP] disallow if exact same cve + scope request already exists even if for deferral",
			existingReqs: []*storage.VulnerabilityRequest{globalDefReq},
			reqType:      storage.VulnerabilityState_FALSE_POSITIVE,
			allow:        false,
		},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			datastore.EXPECT().SearchRawRequests(allAllowedCtx, utils.GetEnforcedRequestsV1Query(cve)).Return(c.existingReqs, nil)
			if c.allow {
				datastore.EXPECT().AddRequest(allAllowedCtx, gomock.Any()).Return(nil)
			}

			var req *storage.VulnerabilityRequest
			if c.reqType == storage.VulnerabilityState_DEFERRED {
				req = fixtures.GetGlobalDeferralRequest(cve)
			} else {
				req = fixtures.GetGlobalFPRequest(cve)
			}
			req.Id = ""
			req.Comments = append(req.Comments, &storage.RequestComment{Message: "comment"})
			err := manager.Create(allAllowedCtx, req)
			fmt.Println(err)
			if c.allow {
				assert.NoError(t, err)
				assert.NotEmpty(t, req.GetId())
			} else {
				assert.Error(t, err)
			}
		})
	}
}

func TestApprovalWithUnifiedDeferral(t *testing.T) {
	testApproval(t)
}

func testApproval(t *testing.T) {
	mockCtrl := gomock.NewController(t)

	testDB := pgtest.ForT(t)

	pendingReqCache, activeReqCache := vulnReqCache.New(), vulnReqCache.New()
	datastore, err := vulReqDS.GetTestPostgresDataStore(t, testDB, pendingReqCache, activeReqCache)
	assert.NoError(t, err)

	imageDataStore := imageDSMocks.NewMockDataStore(mockCtrl)
	imageV2DataStore := imageV2DSMocks.NewMockDataStore(mockCtrl)
	sensorConnMgrMocks := sensorConnMgrMocks.NewMockManager(mockCtrl)
	reprocessor := reprocessorMocks.NewMockLoop(mockCtrl)
	imageCVEDataStore := imageCVEDSMocks.NewMockDataStore(mockCtrl)
	manager := New(nil, datastore, pendingReqCache, activeReqCache, imageDataStore, imageV2DataStore, imageCVEDataStore, sensorConnMgrMocks, reprocessor, nil)

	globalCVE1DefReq := fixtures.GetGlobalDeferralRequest("cve-1")
	globalCVE1FPReq := fixtures.GetGlobalFPRequest("cve-1")
	imageScopedCVE1Req := fixtures.GetImageScopeDeferralRequest("reg1", "img1", "1.0", "cve-1")
	allTagCVE1Req := fixtures.GetImageScopeDeferralRequest("reg1", "img1", ".*", "cve-1")
	otherImageReq := fixtures.GetImageScopeDeferralRequest("reg2", "img1", "1.0", "cve-2")

	globalCVE1DefUpdateReq := fixtures.GetGlobalDeferralRequest("cve-2")
	globalCVE1DefUpdateReq.UpdatedReq = &storage.VulnerabilityRequest_DeferralUpdate{
		DeferralUpdate: &storage.DeferralUpdate{
			CVEs: []string{"cve-1"},
			Expiry: &storage.RequestExpiry{
				Expiry: &storage.RequestExpiry_ExpiresOn{ExpiresOn: protocompat.TimestampNow()},
			},
		},
	}
	globalCVE1DefUpdateReq.Status = storage.RequestStatus_APPROVED_PENDING_UPDATE

	globalCVE1FpUpdateReq := fixtures.GetGlobalFPRequest("cve-2")
	globalCVE1FpUpdateReq.UpdatedReq = &storage.VulnerabilityRequest_FalsePositiveUpdate{
		FalsePositiveUpdate: &storage.FalsePositiveUpdate{
			CVEs: []string{"cve-1"},
		},
	}
	globalCVE1FpUpdateReq.Status = storage.RequestStatus_APPROVED_PENDING_UPDATE

	existingReqs := []*storage.VulnerabilityRequest{
		globalCVE1DefReq,
		globalCVE1FPReq,
		imageScopedCVE1Req,
		allTagCVE1Req,
		otherImageReq,
		globalCVE1DefUpdateReq,
		globalCVE1FpUpdateReq,
	}

	approver := &storage.SlimUser{
		Id:   "approver",
		Name: "approver",
	}
	mockID := mockIdentity.NewMockIdentity(mockCtrl)
	approverCtx := authn.ContextWithIdentity(sac.WithAllAccess(context.Background()), mockID, t)
	deniedReqQ := search.NewQueryBuilder().AddExactMatches(search.RequestStatus, storage.RequestStatus_DENIED.String()).ProtoQuery()
	for _, tc := range []struct {
		desc                    string
		approvedReq             *storage.VulnerabilityRequest
		expectedDeclined        []string
		expectedDeniedUpdateIDs []string
	}{
		{
			desc:                    "decline cve-1 requests and update requests in covered scopes; global, all tags, specific tag",
			approvedReq:             globalCVE1DefReq,
			expectedDeclined:        []string{globalCVE1FPReq.GetId(), allTagCVE1Req.GetId(), imageScopedCVE1Req.GetId()},
			expectedDeniedUpdateIDs: []string{globalCVE1DefUpdateReq.GetId(), globalCVE1FpUpdateReq.GetId()},
		},
		{
			desc:             "decline cve-1 requests in covered scopes; all tags, specific tag",
			approvedReq:      allTagCVE1Req,
			expectedDeclined: []string{imageScopedCVE1Req.GetId()},
		},
		{
			desc:        "none declined",
			approvedReq: otherImageReq,
		},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			for _, req := range existingReqs {
				assert.NoError(t, datastore.AddRequest(allAccessCtx, req))
			}

			mockID.EXPECT().UID().Return(approver.GetId()).AnyTimes()
			mockID.EXPECT().FullName().Return(approver.GetName()).AnyTimes()
			mockID.EXPECT().FriendlyName().Return(approver.GetName()).AnyTimes()

			if features.FlattenImageData.Enabled() {
				imageV2DataStore.EXPECT().Search(gomock.Any(), gomock.Any()).Return(nil, nil)
			} else {
				imageDataStore.EXPECT().Search(gomock.Any(), gomock.Any()).Return(nil, nil)
			}
			_, err := manager.Approve(approverCtx, tc.approvedReq.GetId(), &common.VulnRequestParams{
				Comment: "test approval",
			})
			assert.NoError(t, err)
			results, err := datastore.Search(allAccessCtx, deniedReqQ)
			assert.NoError(t, err)
			assert.ElementsMatch(t, tc.expectedDeclined, search.ResultsToIDs(results))

			deniedUpdateIDs := getDeniedUpdateReqIDs(t, datastore)
			assert.ElementsMatch(t, tc.expectedDeniedUpdateIDs, deniedUpdateIDs)

			for _, req := range existingReqs {
				assert.NoError(t, datastore.RemoveRequestsInternal(allAccessCtx, []string{req.GetId()}))
			}
		})
	}
}

func getDeniedUpdateReqIDs(t *testing.T, datastore vulReqDS.DataStore) []string {
	// When a req with status RequestStatus_APPROVED_PENDING_UPDATE is denied, its status is set to RequestStatus_APPROVED
	// and the comment gives the explanation of denial
	approvedReqQ := search.NewQueryBuilder().AddExactMatches(search.RequestStatus, storage.RequestStatus_APPROVED.String()).ProtoQuery()
	reqs, err := datastore.SearchRawRequests(allAccessCtx, approvedReqQ)
	assert.NoError(t, err)

	var deniedUpdateIDs []string
	for _, req := range reqs {
		if len(req.GetComments()) > 0 && strings.Contains(req.GetComments()[len(req.GetComments())-1].GetMessage(), "declined") {
			deniedUpdateIDs = append(deniedUpdateIDs, req.GetId())
		}
	}
	return deniedUpdateIDs
}
